feat(gateway): upgrade to wavekv 2.0 delta-state sync with v1 dual-stack - #1031
Open
kvinwang wants to merge 21 commits into
Open
feat(gateway): upgrade to wavekv 2.0 delta-state sync with v1 dual-stack#1031kvinwang wants to merge 21 commits into
kvinwang wants to merge 21 commits into
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR upgrades dstack-gateway’s replicated KV sync layer to wavekv 2.0 (delta-state replication) while remaining compatible with wavekv 1.x peers during rolling upgrades, and adds schema-based admission control plus new sync observability fields exposed via the admin RPC.
Changes:
- Add dual-stack HTTP sync endpoints (
/wavekv/syncv1 +/wavekv/sync2v2) and an opportunistic push route (/wavekv/push) to reduce propagation latency. - Enforce per-store key-shape admission via a new schema policy integrated into wavekv node config.
- Extend admin/RPC status reporting with per-store digests and per-peer negotiated protocol / mismatch telemetry, and update wavekv dependency to the v2 branch.
Reviewed changes
Copilot reviewed 9 out of 10 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| dstack/gateway/src/web_routes/wavekv_sync.rs | Adds v2 sync and push HTTP endpoints; refactors gzip handling and introduces envelope decoding. |
| dstack/gateway/src/web_routes.rs | Mounts the new wavekv v2 sync + push routes alongside v1. |
| dstack/gateway/src/kv/sync_service.rs | Extends the sync network interface to use wavekv v2 envelopes and probing for v1/v2 negotiation. |
| dstack/gateway/src/kv/schema.rs | Introduces per-store key admission policy (schema) with tests. |
| dstack/gateway/src/kv/mod.rs | Wires admission policy into wavekv node configs; adds gateway-level wire-compat tests for v1/v2 sync. |
| dstack/gateway/src/kv/https_client.rs | Adds raw-bytes probe POST helper for v2 negotiation and opportunistic push transport. |
| dstack/gateway/src/admin_service.rs | Plumbs new wavekv v2 telemetry (digest, merged/rejected, per-peer protocol/mismatches) into admin RPC responses. |
| dstack/gateway/rpc/proto/gateway_rpc.proto | Extends sync status protos with digest + v2 peer telemetry; deprecates buffered_logs. |
| dstack/Cargo.toml | Switches wavekv dependency to the v2 git branch (with TODO to repoint to crates.io 2.0). |
| dstack/Cargo.lock | Locks wavekv to the v2 git revision and updates transitive deps accordingly. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+240
to
+264
| cert: Option<Certificate<'_>>, | ||
| store: &str, | ||
| data: Data<'_>, | ||
| ) -> Result<Status, Status> { | ||
| verify_gateway_peer(state, cert)?; | ||
|
|
||
| let Some(ref wavekv_sync) = state.wavekv_sync else { | ||
| return Err(Status::ServiceUnavailable); | ||
| }; | ||
|
|
||
| let env = read_envelope(data).await?; | ||
| if env.sender_id == 0 { | ||
| warn!("rejected push from invalid node_id 0"); | ||
| return Err(Status::BadRequest); | ||
| } | ||
|
|
||
| let Some(result) = wavekv_sync.handle_push(store, env) else { | ||
| return Err(Status::NotFound); | ||
| }; | ||
| result.map_err(|e| { | ||
| tracing::error!("{store} push failed: {e:#}"); | ||
| Status::InternalServerError | ||
| })?; | ||
| Ok(Status::Ok) | ||
| } |
Comment on lines
371
to
375
| /// Encode a KV value as MessagePack. | ||
| /// | ||
| /// Structs are encoded as maps keyed by field name rather than as positional | ||
| /// arrays. Field-name keys let a reader skip fields it does not know and fill | ||
| /// in `#[serde(default)]` fields it does not receive, so the value types below |
Comment on lines
+87
to
+92
| let bytes = data | ||
| .open(16.mebibytes()) | ||
| .into_bytes() | ||
| .await | ||
| .map_err(|_| Status::BadRequest)?; | ||
| let decompressed = gunzip(&bytes)?; |
Pick up the wavekv fix for the opportunistic push envelope, which was built without a `sender_uuid` and so failed `check_uuid` on every push — this gateway implements `query_uuid`, so the push channel never worked here. Writes still converged over the periodic round, but each one waited a full sync interval instead of the coalesce window and the receiver logged an error per push blaming node-id reuse. That fix also widens `link_status` to report every known peer rather than only those in the link cache. A peer whose rounds all fail was previously absent from `WaveKvStatus` entirely: a 5xx deliberately does not demote a peer to "v1", so nothing about it moved. Report the new `consecutive_failures` streak so that stall is visible. Document the one direction in which the store schema is not forward compatible: values may gain fields freely, but a new *key* is rejected by nodes that predate it, and a rejection parks ack adoption for the whole round (rule R1). The pair then re-exchanges the same batch indefinitely with no error. New keys therefore ship in two releases — widen the schema everywhere first, write the key second. Also silence a `manual_repeat_n` lint in the pp tests, unrelated but newly raised by the toolchain and enough to fail `clippy -D warnings`.
The HTTP layer was the one part of the sync path with no coverage. It was skipped on the grounds that constructing a `WaveKvSyncService` needs real TLS material; that was wrong. `rcgen` is already a dependency and already used by the cert_store tests, and `verify_gateway_peer` short-circuits under `insecure_skip_attestation`, so a self-signed CA plus a leaf written to a TempDir is enough to build a serving gateway. What this pins that nothing else did: - 503, not 404, when sync is disabled. 404 is the negotiation signal, so a sync-disabled node answering 404 would be cached as "v1" by every peer for a whole reprobe window — and sync is off, so nothing would correct it. - 404 for an unknown store, which is the same signal used deliberately. - An unstamped push is refused at the route and writes nothing. This is the server-side view of the envelope-identity bug; the sender-side view lives in the wavekv push test. - A well-formed push reaches the store, a v2 round trip returns a decodable envelope, and node id 0 is refused. Also stop reporting a 404 on the push route as a delivered push. `post_bytes_probe` maps 404/405 to `Ok(None)` so the v2 probe can read it as "not upgraded yet", but `push_to` discarded the `Option`. A mistyped push URL was therefore indistinguishable from success — the same shape of silent failure that let the unstamped-envelope bug survive, since pushes are best-effort and only debug-logged.
Takes the wavekv fix that verifies the responder's uuid on a v2 response. The field was already on the wire and populated by the responder; only the initiator never read it, so node-id-reuse detection ran in one direction.
The responder-side identity check shipped in the previous bump wedged any peer that regenerated its uuid — an ordinary CVM rebuild, since the uuid is derived from the data directory while the node id comes from config.
The sync wire is gzipped and the 16 MiB cap on a request body caps the *compressed* size, which bounds nothing on its own — gzip expands by three orders of magnitude on attacker-chosen input, so that cap admits a payload that expands into the gigabytes. Every gateway in a cluster shares one app_id, so mTLS proves only that the sender is some gateway of this deployment; it is the same trust level the key schema already treats as insufficient. All four decompression points are now bounded through one helper: both server routes and both client response paths. The client also read peer responses with `Body::collect`, which has no limit at all, so the memory was already spent before any decoding bound could apply; response bodies now go through `Limited` with the same 16 MiB the routes accept on a request. The decompressed ceiling is 128 MiB, far above any legitimate payload: a v2 delta is capped by `max_delta_bytes` at 4 MiB, and the v1 shim answers with the whole live state, which is bounded by the gateway's own key set rather than by anything a peer controls. Tested at the limit as well as past it — a fixture landing exactly on the ceiling must still decode, or the bound could tighten by a byte with only the bomb test still passing.
A local write allocates a sequence number. After a data-directory loss the node keeps its id but has no record of which numbers it already spent — only its peers do — so `bootstrap` rebuilds the counter from their coverage. Anything written before that reuses numbers the peers already treat as seen, and peers filter those writes out of every delta with no error on either side. The three records written at startup were exactly the ones that must not be dropped: `node/info` carries the fresh uuid peers check us against, and `__peer_addr` carries the address they route to. A rebuilt gateway therefore wedged in both directions and stayed wedged until per-peer digest repair happened to fire. They could not simply be moved, because `HttpSyncNetwork::new` read this node's uuid back out of the store, making the `node/info` write a prerequisite of building the sync service at all. That read is the actual defect: our own uuid is local configuration, not replicated state, and routing it through the store created the ordering constraint that forced the bug. It is now passed in, and all three writes happen after the bootstrap. Also bound the response body in `post_json`. Every other response is read through `read_body_bounded`; this one collected without a limit, so a peer could stream until memory ran out. It is the bootnode GetPeers path, and the threat model does not assume a bootnode is honest.
Picks up: WAL truncation of a damaged tail before appending (writes after a torn-tail recovery were silently lost on the next restart), the reset_acks hint the divergence repair always needed (repair reached only entries the peer itself authored), sequence-number recovery that survives an own entry losing LWW, cross-page R1 enforcement, and requests no longer disclosing our state digest. The wire test framed a *request* to check the digest survives transport. Requests no longer carry one, so it frames a response — the direction the digest actually travels — and asserts the request has none.
…ity fixes Retiring a peer no longer discards our coverage of the entries it authored, which is what let the rest of the cluster compute a GC watermark for that origin; membership ops are now WAL-durable, so a peer lost to a crash can no longer widen the watermark and resurrect that peer's deletes.
Mutation testing found `verify_gateway_peer` replaceable with `Ok(())` without turning the suite red. The sync routes are the cluster's write surface — anything reaching them inserts entries that replicate to every gateway — and that function is the only thing in front of them. The cause was in the fixture: every route test sets `insecure_skip_attestation`, which is the function's first statement, so no test had ever executed a line of the gate. The comment claimed the flag "stands in for the mTLS peer check". It does not stand in for it; it removes it. Two gaps, so two changes. `enforcing_gateway` runs with the bypass off. Rocket's local client speaks no TLS and so presents no certificate, which is exactly the case that must be refused, and all three routes are asserted to answer 401. The app-id comparison needed a certificate, and `rocket::mtls::Certificate` has no public constructor — it exists only as the output of a real handshake. But the adapter over it only ever used `cert.extensions()`, which is public and whose element type comes straight out of `X509Certificate`. `RocketCert` now holds the extension list, so a test can build one from a certificate minted in process, and the authorization rule is split out from the Rocket plumbing it was tangled with. None of this needs a TEE or a simulator: the check reads two X.509 extensions and compares bytes. `CertRequest` adds `PHALA_RATLS_APP_ID` unconditionally, and the gateway's own app id is already a constructor parameter. Four cases now pinned: matching id accepted, foreign id forbidden, certificate without an app id refused, and a gateway with no app id of its own authorizing nobody. Each was verified to die under the mutation it targets.
The v1 route had no round-trip test. Mutation testing could delete either store arm, invert the node-id-zero guard, or replace the whole response body with three bytes, and the suite stayed green — every route test targeted v2, because v1 is the compatibility path and attention went to the new one. Deleting the `"persistent"` arm is the sharpest of these. It falls through to `_ => 404`, and a 404 on a sync route is exactly the signal a v2 peer reads as "this node has no such route" — so a broken store dispatch would not surface as an error, it would surface as a successful protocol downgrade, cluster-wide and silently, for a whole reprobe window. The suite already documents that reasoning for the sync-disabled 503 case; the v1 route just had nothing enforcing it. Three tests: a round trip that asserts the response decodes and carries the state this node holds, the same for the ephemeral store, and a node-id-zero rejection matching the push and v2 routes.
…d the key schema Three gaps mutation testing found, none of which needed any infrastructure — all three are pure functions over bytes. `AppIdValidator::validate` had no tests at all. It runs during the TLS handshake, so a validator that always returns `Ok(())` means this gateway completes a mutually-authenticated connection to any peer holding any certificate our CA signed, and then sends it our state. It is the client-side mirror of the route check covered in 8d04ff2, and it was equally undefended: replacing the body with `Ok(())` or inverting the comparison left the suite green. The decompression-limit test asserted a payload of exactly `MAX_DECOMPRESSED_SYNC_BYTES` is accepted — building that payload from the same constant. It therefore held for whatever the constant said, and shrinking 128 MiB to a few kilobytes kept it green while rejecting every real delta. It pinned `>` against `>=` and nothing else. The limits are now checked against what the protocol actually produces: room for wavekv's 4 MiB delta cap, and a compressed ceiling equal to what the routes accept on a request. The key namespace had no tests either. Every builder and parser survived mutation: `handshake_prefix` could return `""`, `parse_inst_key` could return `Some("xyzzy")`. These strings are how a gateway finds its own state after an upgrade, so changing one orphans every existing record — still replicated, still in the digest, unreachable by any reader. Four properties are now pinned: a prefix matches the keys it iterates, a prefix does not capture a neighbour (`inst-a` must not swallow `inst-ab`), builders and parsers round-trip, and a parser refuses a key from another namespace.
A node rebuilt from an empty data directory no longer adopts a requester's ack map while it is bootstrapping — the window in which its own coverage is unknown to it, and in which adopting would have it claim coverage of state it does not hold.
`post_bytes_probe`'s mapping of 404/405 to `Ok(None)` *is* the v1/v2 negotiation: a gateway that has not been upgraded has no `/wavekv/sync2` route, and that status is the only signal its peers get. Every mutation of the condition survived — `||` to `&&`, either `==` to `!=`, the `!` on `is_success` — because nothing exercised the function at all. It cannot be reached without a peer that speaks TLS, since the client is built `https_only()`, and that was enough friction for the whole file to sit at zero. A listener on 127.0.0.1 with a certificate minted in process is enough. No container, no simulator: nothing on this path verifies a quote. Four cases: both "no such route" statuses read as not-upgraded; a 5xx or 4xx stays an error, because reading one as not-upgraded would demote a healthy v2 peer to the v1 path for a whole reprobe window; a 200 is decompressed and returned; and an oversized body is refused. The last one initially passed with the bound removed entirely. Its payload was not valid gzip, so `gunzip_bounded` rejected it whatever the ceiling said, and the assertion measured nothing. It now sends stored-mode gzip — valid, and large enough to clear the compressed ceiling while decompressing well inside the decompressed one — so only the bound under test can reject it.
Three survivors were left on the TLS client after the negotiation tests, and all three sit on paths this cluster depends on. `post_compressed_msg` is the v1 sync path — how a v2 gateway talks to one that has not been upgraded. Its status check was as untested as the negotiation's, so a v1 peer answering 500 could have been decoded as a successful round. `post_json` is the bootnode GetPeers path, where the threat model does not assume the peer is honest, and a failure status must not be parsed as a peer list. The third needed the custom-verifier path. `AppIdValidator` runs inside `CustomCertVerifier`, which rustls only reaches once standard chain verification passes, so unit-testing the validator alone leaves the wiring between them untested — and the wiring is what decides whether a peer from another app can open a connection at all. It now serves a certificate carrying a foreign app id and asserts the handshake fails before any application bytes move, with the matching id as the control. Deleting the validator call turns it red.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Upgrades dstack-gateway onto wavekv 2.0 (delta-state replication, Phala-Network/wavekv#3) as a dual-stack node: it serves the native v2 protocol and keeps serving v1 peers, so a gateway cluster can be upgraded one CVM at a time.
Problem
The gateway's replicated state is an LWW CRDT, but wavekv 1.x replicates it with per-origin operation logs. Op-based replication needs exactly-once ordered delivery, and the machinery that buys it is where the gateway's real failure modes live:
local_ack/peer_acktrack log positions, not state. Two gateways whose WireGuard peer sets have drifted apart report identical healthy status, and the log that would repair them has been truncated. There is no metric an operator can alarm on.apply_pushed_entriesdiscards a whole batch when the first entry's seq is ahead oflocal_ack + 1, and leaves recovery to the pull path noticing later.gateway.tomlshipsinterval = "1m"), so an instance registered on node A is unroutable through node B for up to a minute.Fix
Dual-stack sync
HttpSyncNetworkgains a v2 leg posting to/wavekv/sync2/{store}. A peer still on wavekv 1.x has no such route and answers 404, whichpost_bytes_probesurfaces asOk(None)— distinct from a transport error. wavekv'sSyncManagerreads that as "this peer is v1", falls back to/wavekv/sync, caches the verdict per peer, and re-probes everyprotocol_reprobeso an upgraded peer is picked up without a restart.Serving the other direction needs nothing beyond mounting the route: a v2 gateway answers v1 peers through wavekv's compatibility shim, whose
is_snapshot = trueresponse makes an unmodified v1 client adopt coverage and merge in exactly delta-state order./wavekv/push/{store}carries opportunistic pushes. Per wavekv's rule R3 these merge data only and never move ack coverage, so loss, duplication and reordering are harmless and the periodic round stays the anti-entropy backstop. This is what cuts propagation latency from the sync interval to the coalescing window.Both new routes reuse
verify_gateway_peer(same-app_id mTLS) and the 16 MiB body cap, and decode throughSyncEnvelope::decode, which enforces the schema version and rejects trailing bytes — deliberately not the genericdecodeused for KV values.Admission control
kv/schema.rsconfines each store to the key shapes the gateway actually defines. wavekv enforces it insidemerge, which is the only place covering both sync directions — a check in the HTTP handler would see inbound requests but not entries arriving in a response. A rejection also parks that round's ack adoption (rule R1), so a peer sending inadmissible data keeps re-offering it rather than having it silently dropped.The two stores have disjoint schemas, so an ephemeral-store peer cannot plant
cert/...orinst/...keys.Observability
WaveKvStatusnow reports, per store, the state digest (hex SHA-256 over the replicated state) plus merged/rejected counters, and per peer the negotiatedprotocol("v1"/"v2"),heard_from, anddigest_mismatches.The digest is the operational point of this whole change: two converged replicas produce equal digests by construction, so comparing them across the cluster is both the promotion gate for the rollout and the standing divergence check afterwards.
buffered_logsis kept and marked deprecated — it is always 0 now — so existing clients keep decoding.Verification
cargo test -p dstack-gateway: 90 pass. The cross-version behaviour itself is covered exhaustively in Phala-Network/wavekv#3, whose suite runs the real, unmodified wavekv 1.0 crate from crates.io against v2 (mixed clusters, shim adoption, rollback, tombstones across versions, fault injection, clock skew). This PR adds the gateway-layer wire tests that suite cannot see:a_positionally_encoded_v1_request_is_still_acceptedSyncMessageencoded by a wavekv 1.x gateway still decodes herea_v1_peer_can_decode_our_sync_responsea_v2_envelope_survives_the_transport_framingthe_v1_shim_serves_a_complete_deltais_snapshot = truemerged_entries_outside_the_schema_are_refusedkv::schema(3 tests)cargo fmt --all --checkclean; clippy clean apart from the pre-existingmanual_repeat_ningateway/src/pp.rs:254.Rollout
Per wavekv RFC 0001 §8.4, upgrade one gateway CVM at a time. After each node, the promotion gate is cluster-wide digest equality via
WaveKvStatus, plusprotocolflipping to"v2"for upgraded pairs anddigest_mismatchesstaying at 0. Any anomaly: roll that node back alone — v2 writes the v1 snapshot container and a WAL that is a strict subset of the v1 op set, so a v1 binary reads the same data directory.Note that a pre-existing divergence in a live cluster will surface as a digest mismatch during the rollout. That is the tool working as intended — wavekv 1.x could not have told you — but operators should expect it rather than read it as an upgrade regression.
Follow-up
dstack/Cargo.tomlpoints wavekv at the PR branch. It must be repointed towavekv = "2.0"once Phala-Network/wavekv#3 is merged and released; the TODO is inline. This PR should not merge before that.